{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W2D4_OptimalControl/W2D4_Tutorial1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Neuromatch Academy: Week 3, Day 3, Tutorial 1\n",
    "# Optimal Control for Discrete States and Actions\n",
    "\n",
    "__Content creators:__ Zhengwei Wu, Shreya Saxena, Xaq Pitkow\n",
    "\n",
    "__Content reviewers:__ Karolina Stosio, Roozbeh Farhoodi, Saeed Salehi, Spiros Chavlis, Matt Krause and Michael Waskom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Tutorial Objectives\n",
    "\n",
    "In this tutorial, we will implement a binary control task: a Partially Observable Markov Decision Process (POMDP) that describes fishing. The agent (you) seeks reward from two fishing sites without directly observing where the school of fish is (a group of fish is called a school!). This makes the world a Hidden Markov Model. Based on when and where you catch fish, you keep updating your belief about the fish location, _i.e._ the posterior of the fish given past observations. You should control your position to get the most fish while minimizing the cost of switching sides.\n",
    "\n",
    "You've already learned about stochastic dynamics, latent states, and measurements. Now we introduce you to the new concepts of **control, utility, and policy**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "## Setup  \n",
    "Please execute the cells below to initialize the notebook environment.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.095711Z",
     "iopub.status.busy": "2021-05-25T01:18:16.095165Z",
     "iopub.status.idle": "2021-05-25T01:18:16.472564Z",
     "shell.execute_reply": "2021-05-25T01:18:16.471541Z"
    }
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "import numpy as np\n",
    "import scipy\n",
    "from scipy.linalg import inv\n",
    "from math import isclose\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.488110Z",
     "iopub.status.busy": "2021-05-25T01:18:16.479375Z",
     "iopub.status.idle": "2021-05-25T01:18:16.573856Z",
     "shell.execute_reply": "2021-05-25T01:18:16.573324Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Figure Settings\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "import ipywidgets as widgets\n",
    "from ipywidgets import interact, fixed, HBox, Layout, VBox, interactive, Label\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.611480Z",
     "iopub.status.busy": "2021-05-25T01:18:16.587211Z",
     "iopub.status.idle": "2021-05-25T01:18:16.614571Z",
     "shell.execute_reply": "2021-05-25T01:18:16.614970Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Helper functions\n",
    "\n",
    "binomial = np.random.binomial\n",
    "\n",
    "\n",
    "class ExcerciseError(AssertionError):\n",
    "  pass\n",
    "\n",
    "\n",
    "def test_policy_threshold():\n",
    "  well_done = True\n",
    "  for loc in [-1, 1]:\n",
    "    threshold = 0.4\n",
    "    belief = np.array([.2, .3])\n",
    "    if policy_threshold(threshold, belief, loc) != \"switch\":\n",
    "      raise ExcerciseError(\"'policy_threshold' function is not correctly implemented!\")\n",
    "  for loc in [1, -1]:\n",
    "    threshold = 0.6\n",
    "    belief = np.array([.7, .8])\n",
    "    if policy_threshold(threshold, belief, loc) != \"stay\":\n",
    "      raise ExcerciseError(\"'policy_threshold' function is not correctly implemented!\")\n",
    "  print(\"Well Done!\")\n",
    "\n",
    "\n",
    "def test_value_function():\n",
    "  measurement = np.array([0, 0, 0, 1, 0, 0, 0, 0, 1, 1])\n",
    "  act = np.array([\"switch\", \"stay\", \"switch\", \"stay\", \"stay\",\n",
    "                  \"stay\", \"switch\", \"switch\", \"stay\", \"stay\"])\n",
    "  cost_sw = .5\n",
    "  if not isclose(value_function(measurement, act, cost_sw), .1):\n",
    "    raise ExcerciseError(\"'value_function' function is not correctly implemented!\")\n",
    "  print(\"Well Done!\")\n",
    "\n",
    "\n",
    "def plot_fish(fish_state, ax=None):\n",
    "  \"\"\"\n",
    "  Plot the fish dynamics\n",
    "  \"\"\"\n",
    "  T = len(fish_state)\n",
    "\n",
    "  showlen = min(T, 200)\n",
    "  startT = 0\n",
    "\n",
    "  endT = startT + showlen\n",
    "  showT = range(startT, endT)\n",
    "  time_range = np.linspace(0, showlen - 1)\n",
    "\n",
    "  if not ax:\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(12, 2.5))\n",
    "  ax.plot(- fish_state[showT], color='dodgerblue', markersize=10, linewidth=3.0)\n",
    "  ax.set_xlabel('time', fontsize=18)\n",
    "  ax.set_ylabel('Fish state', rotation=360, fontsize=18)\n",
    "  ax.yaxis.set_label_coords(-0.1, 0.25)\n",
    "  ax.set_xticks([0, showlen, showlen])\n",
    "  ax.tick_params(axis='both', which='major', labelsize=18)\n",
    "  ax.set_xlim([0, showlen])\n",
    "  ax.set_ylim([-1.1, 1.1])\n",
    "  ax.set_yticks([-1, 1])\n",
    "  ax.tick_params(axis='both', which='major', labelsize=18)\n",
    "  labels = [item.get_text() for item in ax.get_yticklabels()]\n",
    "  labels[0] = 'Right'\n",
    "  labels[1] = 'Left'\n",
    "  ax.set_yticklabels(labels)\n",
    "\n",
    "\n",
    "def plot_measurement(measurement, ax=None):\n",
    "  \"\"\"\n",
    "  Plot the measurements\n",
    "  \"\"\"\n",
    "  T = len(measurement)\n",
    "\n",
    "  showlen = min(T, 200)\n",
    "  startT = 0\n",
    "  endT = startT + showlen\n",
    "  showT = range(startT, endT)\n",
    "  time_range = np.linspace(0, showlen - 1)\n",
    "\n",
    "  if not ax:\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(12, 2.5))\n",
    "  ax.plot(measurement[showT], 'r*', markersize=5)\n",
    "  ax.set_xlabel('time', fontsize=18)\n",
    "  ax.set_ylabel('Measurement', rotation=360, fontsize=18)\n",
    "  ax.yaxis.set_label_coords(-0.2, 0.4)\n",
    "  ax.set_xticks([0, showlen, showlen])\n",
    "  ax.tick_params(axis='both', which='major', labelsize=18)\n",
    "  ax.set_xlim([0, showlen])\n",
    "  ax.set_ylim([-.1, 1.1])\n",
    "  ax.set_yticks([0, 1])\n",
    "  ax.set_yticklabels(['no fish', 'caught fish'])\n",
    "\n",
    "\n",
    "def plot_act_loc(loc, act, ax_loc=None):\n",
    "  \"\"\"\n",
    "  Plot the action and location of 200 time points\n",
    "  \"\"\"\n",
    "  T = len(act)\n",
    "\n",
    "  showlen = min(T, 200)\n",
    "  startT = 0\n",
    "\n",
    "  endT = startT + showlen\n",
    "  showT = range(startT, endT)\n",
    "  time_range = np.linspace(0, showlen - 1)\n",
    "\n",
    "  if not ax_loc:\n",
    "    fig, ax_loc = plt.subplots(1, 1, figsize=(12, 2.5))\n",
    "\n",
    "  act_int = (act == \"switch\").astype(int)\n",
    "  ax_loc.plot(-loc[showT], 'g.-', markersize=8, linewidth=5)\n",
    "  ax_loc.plot((act_int[showT] * 4 - 3) * .5, 'rv', markersize=12,\n",
    "              label='switch')\n",
    "  ax_loc.set_xlabel('time', fontsize=18)\n",
    "  ax_loc.set_ylabel('Your state', rotation=360, fontsize=18)\n",
    "\n",
    "  ax_loc.legend(loc=\"upper right\", fontsize=12)\n",
    "  ax_loc.set_xlim([0, showlen])\n",
    "  ax_loc.set_ylim([-1.1, 1.1])\n",
    "  ax_loc.set_yticks([-1, 1])\n",
    "  ax_loc.set_xticks([0, showlen, showlen])\n",
    "  ax_loc.tick_params(axis='both', which='major', labelsize=18)\n",
    "  labels = [item.get_text() for item in ax_loc.get_yticklabels()]\n",
    "  labels[1] = 'Left'\n",
    "  labels[0] = 'Right'\n",
    "  ax_loc.set_yticklabels(labels)\n",
    "\n",
    "\n",
    "def plot_belief(belief, ax1=None, choose_policy=None):\n",
    "  \"\"\"\n",
    "  Plot the belief dynamics of 200 time points\n",
    "  \"\"\"\n",
    "\n",
    "  T = belief.shape[1]\n",
    "\n",
    "  showlen = min(T, 200)\n",
    "  startT = 0\n",
    "\n",
    "  endT = startT + showlen\n",
    "  showT = range(startT, endT)\n",
    "  time_range = np.linspace(0, showlen - 1)\n",
    "\n",
    "  if not ax1:\n",
    "      fig, ax1 = plt.subplots(1, 1, figsize=(12, 2.5))\n",
    "\n",
    "  ax1.plot(belief[0, showT], color='dodgerblue', markersize=10, linewidth=3.0)\n",
    "  ax1.yaxis.set_label_coords(-0.1, 0.25)\n",
    "  ax1.set_xlabel('time', rotation=360, fontsize=18)\n",
    "  ax1.set_ylabel('Belief on \\n left', rotation=360, fontsize=18)\n",
    "  ax1.tick_params(axis='both', which='major', labelsize=18)\n",
    "  ax1.set_xlim([0, showlen])\n",
    "  ax1.set_yticks([0, 1])\n",
    "  ax1.set_ylim([0, 1.1])\n",
    "  ax1.set_xticks([0, showlen, showlen])\n",
    "\n",
    "  if choose_policy == \"threshold\":\n",
    "    ax2 = ax1.twinx()\n",
    "    ax2.plot(time_range, threshold * np.ones(time_range.shape), 'r--')\n",
    "    ax2.plot(time_range, (1 - threshold) * np.ones(time_range.shape), 'c--')\n",
    "    ax2.set_yticks([threshold, 1 - threshold])\n",
    "    ax2.set_ylim([0, 1.1])\n",
    "    ax2.tick_params(axis='both', which='major', labelsize=18)\n",
    "    labels = [item.get_text() for item in ax2.get_yticklabels()]\n",
    "    labels[0] = 'threshold to switch \\n from left to right'\n",
    "    labels[-1] = 'threshold to switch \\n from right to left'\n",
    "    ax2.set_yticklabels(labels)\n",
    "\n",
    "\n",
    "def plot_dynamics(belief, loc, act, meas, fish_state, choose_policy):\n",
    "  \"\"\"\n",
    "  Plot the dynamics of 200 time points\n",
    "  \"\"\"\n",
    "  if choose_policy == 'threshold':\n",
    "    fig, [ax0, ax_loc, ax1, ax_bel] = plt.subplots(4, 1, figsize=(12, 9))\n",
    "    plot_fish(fish_state, ax=ax0)\n",
    "    plot_belief(belief, ax1=ax_bel)\n",
    "    plot_measurement(meas, ax=ax1)\n",
    "    plot_act_loc(loc, act, ax_loc=ax_loc)\n",
    "  else:\n",
    "    fig, [ax0, ax1, ax_bel] = plt.subplots(3, 1, figsize=(12, 7))\n",
    "    plot_fish(fish_state, ax=ax0)\n",
    "    plot_belief(belief, ax1=ax_bel)\n",
    "    plot_measurement(meas, ax=ax1)\n",
    "\n",
    "  plt.tight_layout()\n",
    "  plt.show()\n",
    "\n",
    "\n",
    "def belief_histogram(belief, bins=100):\n",
    "  \"\"\"\n",
    "  Plot the histogram of belief states\n",
    "  \"\"\"\n",
    "  fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
    "  ax.hist(belief, bins)\n",
    "  ax.set_xlabel('belief', fontsize=18)\n",
    "  ax.set_ylabel('count', fontsize=18)\n",
    "  plt.show()\n",
    "\n",
    "\n",
    "def plot_value_threshold(cost_sw=0.5, T=10000, p_stay=.95,\n",
    "                         high_rew_p=.4, low_rew_p=.1, step=.05):\n",
    "  \"\"\"\n",
    "  Helper function to plot the value function and threshold\n",
    "  \"\"\"\n",
    "  params = [T, p_stay, high_rew_p, low_rew_p, _]\n",
    "\n",
    "  threshold_array, value_array = value_threshold(params, cost_sw, step)\n",
    "  yrange = np.max(value_array) - np.min(value_array)\n",
    "\n",
    "  fig_, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
    "  ax.plot(threshold_array, value_array, 'b')\n",
    "  ax.set_ylim([np.min(value_array) - yrange * .1, np.max(value_array) + yrange * .1])\n",
    "  ax.set_title(f'threshold vs value with switching cost c = {cost_sw:.2f}',\n",
    "               fontsize=20)\n",
    "  ax.set_xlabel('threshold', fontsize=16)\n",
    "  ax.set_ylabel('value', fontsize=16)\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 1: Dynamics of Fishing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.624714Z",
     "iopub.status.busy": "2021-05-25T01:18:16.623809Z",
     "iopub.status.idle": "2021-05-25T01:18:16.670703Z",
     "shell.execute_reply": "2021-05-25T01:18:16.671111Z"
    },
    "outputId": "c0da2078-516d-4da7-a50a-e3968b248675"
   },
   "outputs": [],
   "source": [
    "#@title Video 1: Gone fishing\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id='3oIwUFpolVA', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "\n",
    "## Introduction\n",
    "\n",
    "There are two locations for the fish and you (Left and Right). If you're on the same side as the fish, you'll catch more, with probabilty $q_{\\rm high}$ per discrete time step. Otherwise you may still catch fish with probability $q_{\\rm low}$. One fish is worth 1 \"point\".\n",
    "\n",
    "The fish location $s^{\\rm fish}$ is latent. The only information you get about the fish location is when you catch one. Secretly at each time step, the fish may switch sides with a certain probability $p_{\\rm sw} = 1 - p_{\\rm stay}$.\n",
    "\n",
    "\n",
    "You are in control of your own location. You may stay on your current side with no cost, or switch to the other side and incur an action cost $C$ (again, in units of fish).\n",
    "\n",
    "You select controls or actions by following a **policy**. This defines what to do in any situation. Here the situation is specified by your location and your belief $b_t$ about the fish location. For optimal control we assume that this belief is the posterior probability over the current fish location, given all the past measurements. We only need one number for this, since the fish are either on the left or the right. So we write \n",
    "\n",
    "$$b_t = p(s^{\\rm fish}_t = {\\rm Right}\\  |\\  m_{0:t}, a_{0:t-1})$$\n",
    "\n",
    "where $m$ are the measurements, and $a$ are the controls or actions (stay or switch).\n",
    "\n",
    "Ultimately we will parameterize the policy by a simple threshold on beliefs. (This happens to be optimal if you pick the right threshold!) When your belief that fish are on your current side falls below a threshold $\\theta$, you switch to the other side.\n",
    "\n",
    "Your **overall goals** in this tutorial are:\n",
    "1. Measure when fish are caught, first if the school of fish doesn't move.\n",
    "2. For moving fish, plot their dynamics and your belief about it based on your measurements.\n",
    "3. Compute the value for a given control policy.\n",
    "4. Find the optimal policy for controlling your position."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 1.1\n",
    "Evaluate the following methods implemented in `binaryHMM` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.681098Z",
     "iopub.status.busy": "2021-05-25T01:18:16.679260Z",
     "iopub.status.idle": "2021-05-25T01:18:16.683802Z",
     "shell.execute_reply": "2021-05-25T01:18:16.684219Z"
    }
   },
   "outputs": [],
   "source": [
    "# class for the binary HMM task\n",
    "\n",
    "class binaryHMM():\n",
    "\n",
    "  def __init__(self, params, fish_initial=-1, loc_initial=-1):\n",
    "    self.params = params\n",
    "    self.fish_initial = fish_initial\n",
    "    self.loc_initial = loc_initial\n",
    "\n",
    "  def fish_state_telegraph(self, fish_past, p_stay):\n",
    "    \"\"\"\n",
    "    fish state update according to telegraph process\n",
    "\n",
    "    Args:\n",
    "      fish_past (int): the fish location (-1 for left side, 1 for right side)\n",
    "      p_stay : the probability that the state of a certain site stays the same\n",
    "\n",
    "    Returns:\n",
    "      fish_new (int): updated fish location\n",
    "    \"\"\"\n",
    "    # we use logical operation XOR (denoted by ^ in python)\n",
    "    fish_new = (1 - binomial(1, p_stay)) ^ ((fish_past + 1) // 2)\n",
    "    fish_new = fish_new * 2 - 1\n",
    "\n",
    "    return fish_new\n",
    "\n",
    "\n",
    "  def fish_dynamics(self):\n",
    "    \"\"\"\n",
    "    fish state dynamics according to telegraph process\n",
    "\n",
    "    Returns:\n",
    "      fish_state (numpy array of int)\n",
    "    \"\"\"\n",
    "    T, p_stay, _, _, _ = self.params\n",
    "    fish_state = np.zeros(T, int)  # -1: left side ; 1: right side\n",
    "\n",
    "    # initialization\n",
    "    fish_state[0] = self.fish_initial\n",
    "\n",
    "    for t in range(1, T):\n",
    "        fish_state[t] = self.fish_state_telegraph(fish_state[t - 1], p_stay)\n",
    "\n",
    "    return fish_state\n",
    "\n",
    "  def generate_process_lazy(self):\n",
    "    \"\"\"\n",
    "    fish dynamics and measurements if you always stay in the intial location\n",
    "    without changing sides\n",
    "\n",
    "    Returns:\n",
    "      fish_state (numpy array of int): locations of the fish\n",
    "      loc (numpy array of int): left or right site, -1 for left, and 1 for right\n",
    "      measurement (numpy array of binary): whether a reward is obtained\n",
    "    \"\"\"\n",
    "\n",
    "    T, _, high_rew_p, low_rew_p, _ = self.params\n",
    "    rew_p_vector = np.array([low_rew_p, high_rew_p])\n",
    "\n",
    "    fish_state = self.fish_dynamics()\n",
    "    loc  = np.zeros(T, int)            # -1: left side, 1: right side\n",
    "    measurement = np.zeros(T, int)     # 0: no food, 1: get food\n",
    "\n",
    "    for t in range(0, T):\n",
    "      loc[t] = self.loc_initial\n",
    "      # new measurement\n",
    "      measurement[t] = binomial(1, rew_p_vector[(fish_state[t] == loc[t]) * 1])\n",
    "\n",
    "    return fish_state, loc, measurement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 1.2\n",
    "\n",
    "Run the following code to see the dynamics of the fish as they follow the telegraph process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.689755Z",
     "iopub.status.busy": "2021-05-25T01:18:16.689090Z",
     "iopub.status.idle": "2021-05-25T01:18:16.693087Z",
     "shell.execute_reply": "2021-05-25T01:18:16.692585Z"
    }
   },
   "outputs": [],
   "source": [
    "def update_ex_1(p_stay=.95, high_rew_p=.4, low_rew_p=.1, T=200):\n",
    "  \"\"\"\n",
    "    p_stay: probability fish stay\n",
    "    high_rew_p: p(catch fish) when you're on their side\n",
    "    low_rew_p : p(catch fish) when you're on other side\n",
    "  \"\"\"\n",
    "\n",
    "  params = [T, p_stay, high_rew_p, low_rew_p, _]\n",
    "\n",
    "  #### initial condition for fish [fish_initial] and you [loc_initial] ####\n",
    "  binaryHMM_test = binaryHMM(params, fish_initial=-1, loc_initial=-1)\n",
    "\n",
    "  fish_state = binaryHMM_test.fish_dynamics()\n",
    "  plot_fish(fish_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "b6aaf76696e74a9e9e057e5ab6412478",
      "97d83c9eb6b847d88ac4a3ef4cc553c2",
      "0ef4826e52174f33891ad4748675515f",
      "da2a074c5b0e4f649c6c81ac4974e2f1",
      "8eba73aa0b5341179b40d5d725cdcfae",
      "482120bbb75443d5ae4ca7b30505dc01",
      "64af3a773dba4df48641c58c726232a7",
      "63a1056afd414526b0579c20d4638dcc",
      "9afe9152ad86427ca9732176b52099b6",
      "d25840f959fd4091add0815da033bd2e",
      "89731469dfc8423aad89e5d06a0b66ec",
      "f503a1557e6d479b8426807a4f6d0298"
     ]
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.711356Z",
     "iopub.status.busy": "2021-05-25T01:18:16.702149Z",
     "iopub.status.idle": "2021-05-25T01:18:16.745599Z",
     "shell.execute_reply": "2021-05-25T01:18:16.741220Z"
    },
    "outputId": "8eb89404-d8e6-4a95-d9c3-1ee0872379aa"
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "widget=interactive(update_ex_1, {'manual': True},\n",
    "                high_rew_p=fixed(.4),\n",
    "                low_rew_p=fixed(.1),\n",
    "                p_stay=(.5, 1., .001),\n",
    "                T=fixed(200))\n",
    "\n",
    "widget.children[-2].description='Run Simulation'\n",
    "widget.children[-2].style.button_color='lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 2: Catch some fish on each side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.750929Z",
     "iopub.status.busy": "2021-05-25T01:18:16.750366Z",
     "iopub.status.idle": "2021-05-25T01:18:16.797247Z",
     "shell.execute_reply": "2021-05-25T01:18:16.796713Z"
    },
    "outputId": "fd82ef26-d1c1-4420-91e5-135ed9757882"
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Catch some fish\n",
    "video = YouTubeVideo(id='1-Wionllt9U', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 2\n",
    "\n",
    "Now set $p_{\\rm stay} = 1$ so that the state of the two sites are fixed, and we can directly see the chances of catching fish on each side. The variable `fish_initial` indicates the initial side of the fish, and `loc_initial` indicates your initial location. They each take value $-1$ for left and $1$ for right.\n",
    "\n",
    "**Instructions:**\n",
    "1. set the two locations (`fish_initial` and `loc_initial`) to be the _same_, and measure when you catch fish.\n",
    "2. set the two locations (`fish_initial` and `loc_initial`) to be the _different_, and measure when you catch fish.\n",
    "3. visually compare the measurements from 1 and 2.\n",
    "4. Finally, you can also play around with `high_rew_p` (high reward probability) and `low_rew_p` (low reward probability) sliders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.803187Z",
     "iopub.status.busy": "2021-05-25T01:18:16.801948Z",
     "iopub.status.idle": "2021-05-25T01:18:16.803866Z",
     "shell.execute_reply": "2021-05-25T01:18:16.804325Z"
    }
   },
   "outputs": [],
   "source": [
    "def update_ex_2(p_stay=1., high_rew_p=.6, low_rew_p=.05, T=100):\n",
    "  \"\"\"\n",
    "    p_stay: probability fish stay\n",
    "    high_rew_p: p(catch fish) when you're on their side\n",
    "    low_rew_p : p(catch fish) when you're on other side\n",
    "  \"\"\"\n",
    "  params = [T, p_stay, high_rew_p, low_rew_p, _]\n",
    "\n",
    "  #### initial condition for fish [fish_initial] and you [loc_initial] ####\n",
    "  binaryHMM_test = binaryHMM(params, fish_initial=-1, loc_initial=-1)\n",
    "\n",
    "  fish_state, loc, measurement = binaryHMM_test.generate_process_lazy()\n",
    "  plot_measurement(measurement)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "a38428817fc4485d8862808a3f06bb1e",
      "22826f24d1f542e78eac00d4f5d4966d",
      "6a4317ea4e8c431faea40f21c6005fb2",
      "774d5399b1114f99b5a0b429d53b089e",
      "ad7ff1c00d4348efbc51680873a57160",
      "fef4d5169b1043daa4fe83d8c58c150f",
      "cabeadedcf2044238fecc9599d1d7009",
      "aa9eaaa94910456ba0d0908b939ce281",
      "9a5013e0305a4fd682d14aaf7910c7fe",
      "152b7a4951924060bea65872c66cc4e7",
      "4d5251db9b5d4a56bece24ff06cdde35",
      "74f7de7a710b498d9493027bb9274c9e",
      "8d4045d3dbe5430fa4787cf03548f2a7",
      "971b56b03c234b4a9557290fed04c3e5",
      "c01129e6a2a64fa88cf003fcc1c023c4"
     ]
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.853237Z",
     "iopub.status.busy": "2021-05-25T01:18:16.812310Z",
     "iopub.status.idle": "2021-05-25T01:18:16.868404Z",
     "shell.execute_reply": "2021-05-25T01:18:16.860017Z"
    },
    "outputId": "4bb9d690-6515-4c3c-b1f1-ad8fe1b7ed1a"
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "widget=interactive(update_ex_2, {'manual': True},\n",
    "                high_rew_p=(.0, 1., .001),\n",
    "                low_rew_p=(.0, 1., .001),\n",
    "                p_stay=fixed(1.),\n",
    "                T=fixed(100))\n",
    "\n",
    "widget.children[-2].description='Run Simulation'\n",
    "widget.children[-2].style.button_color='lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 3: Belief dynamics and belief distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.873814Z",
     "iopub.status.busy": "2021-05-25T01:18:16.873305Z",
     "iopub.status.idle": "2021-05-25T01:18:16.916505Z",
     "shell.execute_reply": "2021-05-25T01:18:16.915910Z"
    },
    "outputId": "560de3be-f986-47af-993e-918709500bcc"
   },
   "outputs": [],
   "source": [
    "#@title Video 3: Where are the fish?\n",
    "video = YouTubeVideo(id='wCzVnnd4bmg', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 3.1: Plot belief dynamics and belief distributions\n",
    "\n",
    "We have provided a class for the binary task, with the agent always staying at one side following a lazy policy function `def policy_lazy(belief, loc)` that we provided. Now in this exercise, you will extend the module to generate the real dynamics, including beliefs and a moving agent. With the generated data, we will see how beliefs change over time, and how often different beliefs happen.\n",
    "\n",
    "For convenience, your belief at time *t* is actually a 2-dimensional vector. The first element is the belief that the fish are on the left, and the second element is the belief the fish are on the right. At every time, these elements sum to $1$.\n",
    "\n",
    "We will first check the dynamics with lazy policy, and then explore the case with a threshold-based policy.\n",
    "\n",
    "**Instructions:**\n",
    "Evaluate the cells below to setup the lazy policy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.930104Z",
     "iopub.status.busy": "2021-05-25T01:18:16.924025Z",
     "iopub.status.idle": "2021-05-25T01:18:16.932310Z",
     "shell.execute_reply": "2021-05-25T01:18:16.931838Z"
    }
   },
   "outputs": [],
   "source": [
    "class binaryHMM_belief(binaryHMM):\n",
    "\n",
    "  def __init__(self, params, fish_initial = -1, loc_initial = -1,\n",
    "                choose_policy = 'threshold'):\n",
    "    binaryHMM.__init__(self, params, fish_initial, loc_initial)\n",
    "    self.choose_policy = choose_policy\n",
    "\n",
    "  def generate_process(self):\n",
    "    \"\"\"\n",
    "    fish dynamics and measurements based on the choosen policy\n",
    "\n",
    "    Returns:\n",
    "      belief (numpy array of float): belief on the states of the two sites\n",
    "      act (numpy array of string): actions over time\n",
    "      loc (numpy array of int): left or right site\n",
    "      measurement (numpy array of binary): whether a reward is obtained\n",
    "      fish_state (numpy array of int): fish locations\n",
    "    \"\"\"\n",
    "\n",
    "    T, p_stay, high_rew_p, low_rew_p, threshold = self.params\n",
    "    fish_state = self.fish_dynamics()     # -1: left side; 1: right side\n",
    "    loc = np.zeros(T, int)                # -1: left side, 1: right side\n",
    "    measurement = np.zeros(T, int)        # 0: no food, 1: get food\n",
    "    act = np.empty(T, dtype='object')     # \"stay\", or \"switch\"\n",
    "    belief = np.zeros((2, T), float)      # the probability that the fish is on the left (1st element)\n",
    "                                          # or on the right (2nd element),\n",
    "                                          # the beliefs on the two boxes sum up to be 1\n",
    "\n",
    "    rew_prob = np.array([low_rew_p, high_rew_p])\n",
    "\n",
    "    # initialization\n",
    "    loc[0] = -1\n",
    "    measurement[0] = 0\n",
    "    belief_0 = np.random.random(1)[0]\n",
    "    belief[:, 0] = np.array([belief_0, 1 - belief_0])\n",
    "    act[0] = self.policy(threshold, belief[:, 0], loc[0])\n",
    "\n",
    "    for t in range(1, T):\n",
    "      if act[t - 1] == \"stay\":\n",
    "        loc[t] = loc[t - 1]\n",
    "      else:\n",
    "        loc[t] = - loc[t - 1]\n",
    "\n",
    "      # new measurement\n",
    "      measurement[t] = binomial(1, rew_prob[(fish_state[t] == loc[t]) * 1])\n",
    "      belief[0, t] = self.belief_update(belief[0, t - 1] , loc[t],\n",
    "                                        measurement[t], p_stay,\n",
    "                                        high_rew_p, low_rew_p)\n",
    "      belief[1, t] = 1 - belief[0, t]\n",
    "\n",
    "      act[t] = self.policy(threshold, belief[:, t], loc[t])\n",
    "\n",
    "    return belief, loc, act, measurement, fish_state\n",
    "\n",
    "  def policy(self, threshold, belief, loc):\n",
    "    \"\"\"\n",
    "    chooses policy based on whether it is lazy policy\n",
    "        or a threshold-based policy\n",
    "\n",
    "    Args:\n",
    "      threshold (float): the threshold of belief on the current site,\n",
    "          when the belief is lower than the threshold, switch side\n",
    "      belief (numpy array of float): the belief on the two sites\n",
    "      loc (int) : the location of the agent\n",
    "\n",
    "    Returns:\n",
    "      act (string): \"stay\" or \"switch\"\n",
    "    \"\"\"\n",
    "    if self.choose_policy == \"threshold\":\n",
    "      act = policy_threshold(threshold, belief, loc)\n",
    "    if self.choose_policy == \"lazy\":\n",
    "      act = policy_lazy(belief, loc)\n",
    "\n",
    "    return act\n",
    "\n",
    "  def belief_update(self, belief_past, loc, measurement, p_stay,\n",
    "                    high_rew_p, low_rew_p):\n",
    "    \"\"\"\n",
    "    using PAST belief on the LEFT box, CURRENT location and\n",
    "        and measurement to update belief\n",
    "    \"\"\"\n",
    "    rew_prob_matrix = np.array([[1 - high_rew_p, high_rew_p],\n",
    "                                [1 - low_rew_p, low_rew_p]])\n",
    "\n",
    "    # update belief posterior, p(s[t] | measurement(0-t), act(0-t-1))\n",
    "    belief_0 = (belief_past * p_stay  + (1 - belief_past) * (1 - p_stay)) *\\\n",
    "                            rew_prob_matrix[(loc + 1) // 2, measurement]\n",
    "    belief_1 = ((1 - belief_past) * p_stay + belief_past * (1 - p_stay)) *\\\n",
    "                            rew_prob_matrix[1-(loc + 1) // 2, measurement]\n",
    "\n",
    "    belief_0 = belief_0 / (belief_0 + belief_1)\n",
    "\n",
    "    return belief_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.937052Z",
     "iopub.status.busy": "2021-05-25T01:18:16.935895Z",
     "iopub.status.idle": "2021-05-25T01:18:16.937632Z",
     "shell.execute_reply": "2021-05-25T01:18:16.938050Z"
    }
   },
   "outputs": [],
   "source": [
    "def policy_lazy(belief, loc):\n",
    "  \"\"\"\n",
    "  This function is a lazy policy where stay is also taken\n",
    "  \"\"\"\n",
    "  act = \"stay\"\n",
    "\n",
    "  return act"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Excercise 3.1.1 Belief Update (Optional) \n",
    "The belief update function is provided in the `binaryHMM_belief` class in the above excercise. One optional excercise is to write your own code for the belief update based on the past belief `belief_past`, your location `loc`, the observation `measurement`, and the probability that fish stays on one side `p_stay`.\n",
    "\n",
    "You should only do this if you still have time after Exercise 3.2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 3.2: Task dynamics following a **lazy** policy\n",
    "\n",
    "The parameter for policy `choose_policy` can be either \"*lazy*\" or \"*threshold*\". In the following example, use the lazy policy.\n",
    "\n",
    "**Instructions:**\n",
    "* With the class defined above, we have created an object of `binaryHMM_belief` given parameters of the dynamics, *params*, and a parameter for policy.\n",
    "* Run the dynamics and explain the time series of the beliefs you see.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.944200Z",
     "iopub.status.busy": "2021-05-25T01:18:16.943013Z",
     "iopub.status.idle": "2021-05-25T01:18:16.945010Z",
     "shell.execute_reply": "2021-05-25T01:18:16.945467Z"
    }
   },
   "outputs": [],
   "source": [
    "def update_ex_3(p_stay=.98, threshold=.2, high_rew_p=.4, low_rew_p=.1, T=200):\n",
    "  \"\"\"\n",
    "    p_stay: probability fish stay\n",
    "    high_rew_p: p(catch fish) when you're on their side\n",
    "    low_rew_p : p(catch fish) when you're on other side\n",
    "    threshold: threshold of belief below which switching is taken\n",
    "  \"\"\"\n",
    "\n",
    "  params = [T, p_stay, high_rew_p, low_rew_p, threshold]\n",
    "\n",
    "  #### initial condition for fish [fish_initial] and you [loc_initial] ####\n",
    "  binaryHMM_test = binaryHMM_belief(params, choose_policy=\"lazy\",\n",
    "                                    fish_initial=-1, loc_initial=-1)\n",
    "\n",
    "  belief, loc, act, measurement, fish_state = binaryHMM_test.generate_process()\n",
    "  plot_dynamics(belief, loc, act, measurement, fish_state,\n",
    "                binaryHMM_test.choose_policy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "b1e4997dea284637a6d641bd0dc47f06",
      "9c1bec3d582c40408c440c814459e66a",
      "f7d5f61fe10b49088bfaf8381db3a5cc",
      "d2ff36f8fc1d4abfa22b96b087fd9cda",
      "aed0b26916db4d1cb59586e3beed11bd",
      "a562e74700564ef5a8b75a6d3ff11b4b",
      "235568200fac425991721b7d933a5c8a",
      "e11fab3eed404236a7a47c2eb3d0cef9",
      "2125029e45824d1bbbad8ad43fb2bfe9",
      "ecd2fa1da5034d9791f5d87985319d31",
      "6770002b1e074beabb9f9b212e45ddc0",
      "7896306be625455298450006b4a3b293",
      "90140c330ca24df89249cdbc5b6f826a",
      "0f9e6e2854f34a59af48a187f1753458",
      "04e9c82034474b73b88ed19dcce6776c",
      "a24952f457c243fea5e18c9734d6a050",
      "0a046ef4f40646c88f14f3bd78b7e4de",
      "f02e4e46ddd6412b8743d2f88671e70f"
     ]
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:16.998976Z",
     "iopub.status.busy": "2021-05-25T01:18:16.953248Z",
     "iopub.status.idle": "2021-05-25T01:18:17.015903Z",
     "shell.execute_reply": "2021-05-25T01:18:17.003074Z"
    },
    "outputId": "e04c7a1e-94e9-4a43-bdd4-571dd7e07dcb"
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "widget=interactive(update_ex_3, {'manual': True},\n",
    "                high_rew_p=(.0, 1., .001),\n",
    "                low_rew_p=(.0, 1., .001),\n",
    "                p_stay=(.5, 1., .001),\n",
    "                T=fixed(200),\n",
    "                threshold=fixed(.2))\n",
    "\n",
    "widget.children[-2].description='Run Simulation'\n",
    "widget.children[-2].style.button_color='lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 4: Implementing threshold policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.020983Z",
     "iopub.status.busy": "2021-05-25T01:18:17.019769Z",
     "iopub.status.idle": "2021-05-25T01:18:17.067063Z",
     "shell.execute_reply": "2021-05-25T01:18:17.067525Z"
    },
    "outputId": "3d883c89-81d2-4413-cb7e-eb432bf69717"
   },
   "outputs": [],
   "source": [
    "#@title Video 4: How should you act?\n",
    "video = YouTubeVideo(id='G3fNz23IDUg', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 4: dynamics following a **threshold-based** policy.\n",
    "\n",
    "**Instructions:**\n",
    "\n",
    "* You need to code a new policy `def policy_threshold(threshold, belief, loc)`. The policy takes three inputs: your belief about the fish state, your location (\"Left\" or \"Right\"), and a belief _threshold_: when your belief that you are on the same side as the fish drops below this threshold, you choose to switch; otherwise you stay.\n",
    "\n",
    "* You should return an action for each time *t*, which takes the value of \"stay\" or \"switch\".\n",
    "\n",
    "* After you complete the code for the policy based on threshold, create an object of `binaryHMM_belief` and set the policy parameter to be `choose_policy = threshold`.\n",
    "* We have provided an example of the parameters. You should play with the parameters to see the various dynamics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.072909Z",
     "iopub.status.busy": "2021-05-25T01:18:17.071691Z",
     "iopub.status.idle": "2021-05-25T01:18:17.073545Z",
     "shell.execute_reply": "2021-05-25T01:18:17.073982Z"
    }
   },
   "outputs": [],
   "source": [
    "def policy_threshold(threshold, belief, loc):\n",
    "  \"\"\"\n",
    "  chooses whether to switch side based on whether the belief\n",
    "      on the current site drops below the threshold\n",
    "\n",
    "  Args:\n",
    "    threshold (float): the threshold of belief on the current site,\n",
    "                        when the belief is lower than the threshold, switch side\n",
    "    belief (numpy array of float, 2-dimensional): the belief on the\n",
    "                                                  two sites at a certain time\n",
    "    loc (int) : the location of the agent at a certain time\n",
    "                -1 for left side, 1 for right side\n",
    "\n",
    "  Returns:\n",
    "    act (string): \"stay\" or \"switch\"\n",
    "  \"\"\"\n",
    "  ############################################################################\n",
    "  ## Insert code to:\n",
    "  ## generate actions (Stay or Switch) for current belief and location\n",
    "  ##\n",
    "  ## Belief is a 2d vector: first element = Prob(fish on Left | measurements)\n",
    "  ##                       second element = Prob(fish on Right  | measurements)\n",
    "  ## Returns \"switch\" if Belief that fish are in your current location < threshold\n",
    "  ##         \"stay\" otherwise\n",
    "  ##\n",
    "  ## Hint: use loc value to determine which row of belief you need to use\n",
    "  ##       see the docstring for more information about loc\n",
    "  ##\n",
    "  ## complete the function and remove\n",
    "  raise NotImplementedError(\"Student exercise: Please complete <act>\")\n",
    "  ############################################################################\n",
    "  # Write the if statement\n",
    "  if ...:\n",
    "    # action below threshold\n",
    "    act = ...\n",
    "  else:\n",
    "    # action above threshold\n",
    "    act = ...\n",
    "\n",
    "  return act\n",
    "\n",
    "\n",
    "# uncomment the line below to test your function\n",
    "# test_policy_threshold()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.079645Z",
     "iopub.status.busy": "2021-05-25T01:18:17.078462Z",
     "iopub.status.idle": "2021-05-25T01:18:17.081407Z",
     "shell.execute_reply": "2021-05-25T01:18:17.080960Z"
    },
    "outputId": "e336f814-79ab-4e9b-9def-08d9f637d5c5"
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "def policy_threshold(threshold, belief, loc):\n",
    "  \"\"\"\n",
    "  chooses whether to switch side based on whether the belief\n",
    "      on the current site drops below the threshold\n",
    "\n",
    "  Args:\n",
    "    threshold (float): the threshold of belief on the current site,\n",
    "                        when the belief is lower than the threshold, switch side\n",
    "    belief (numpy array of float, 2-dimensional): the belief on the\n",
    "                                                  two sites at a certain time\n",
    "    loc (int) : the location of the agent at a certain time\n",
    "                -1 for left side, 1 for right side\n",
    "\n",
    "  Returns:\n",
    "    act (string): \"stay\" or \"switch\"\n",
    "  \"\"\"\n",
    "  # Write the if statement\n",
    "  if belief[(loc + 1) // 2] <= threshold:\n",
    "    # action below threshold\n",
    "    act = \"switch\"\n",
    "  else:\n",
    "    # action above threshold\n",
    "    act = \"stay\"\n",
    "\n",
    "  return act\n",
    "\n",
    "\n",
    "test_policy_threshold()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.088082Z",
     "iopub.status.busy": "2021-05-25T01:18:17.086926Z",
     "iopub.status.idle": "2021-05-25T01:18:17.088763Z",
     "shell.execute_reply": "2021-05-25T01:18:17.089184Z"
    }
   },
   "outputs": [],
   "source": [
    "def update_ex_4(p_stay=.98, threshold=.2, high_rew_p=.4, low_rew_p=.1, T=200):\n",
    "  \"\"\"\n",
    "    p_stay: probability fish stay\n",
    "    high_rew_p: p(catch fish) when you're on their side\n",
    "    low_rew_p : p(catch fish) when you're on other side\n",
    "    threshold: threshold of belief below which switching is taken\n",
    "  \"\"\"\n",
    "  params = [T, p_stay, high_rew_p, low_rew_p, threshold]\n",
    "\n",
    "  #### initial condition for fish [fish_initial] and you [loc_initial] ####\n",
    "  binaryHMM_test = binaryHMM_belief(params, fish_initial=-1, loc_initial=-1,\n",
    "                                    choose_policy=\"threshold\")\n",
    "\n",
    "  belief, loc, act, measurement, fish_state = binaryHMM_test.generate_process()\n",
    "  plot_dynamics(belief, loc, act, measurement,\n",
    "                fish_state, binaryHMM_test.choose_policy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "2ddab28498ef4be8b8aad1d6ed51936c",
      "fc1acd13b163405287e936b89564abe7",
      "6c76a03ff3d9480f8d57c4269f598aa0",
      "fbacf306283f4cdba7aa4173a79a292e",
      "eff031cbedf0474cb287d4dae0245fa7",
      "4ee72a82b0db452bb6a9f5ab33c07760",
      "1d61a1cd28ad450185d31146c341c7f0",
      "9db2b36861fc4357a167657677907a47",
      "9d8554b1d8654dfe8a9d72f1c4928164",
      "43a95e003c0c4af88b228a14ca866d5f",
      "fdcb1149e66c4c0db2145e7185395be1",
      "cfd0b54f32ce4e189cbfce476bf69fac"
     ]
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.128785Z",
     "iopub.status.busy": "2021-05-25T01:18:17.105182Z",
     "iopub.status.idle": "2021-05-25T01:18:17.143077Z",
     "shell.execute_reply": "2021-05-25T01:18:17.137214Z"
    },
    "outputId": "6de6d976-488d-488b-a6bf-a202446fdee0"
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "widget=interactive(update_ex_4, {'manual': True},\n",
    "                high_rew_p=fixed(.4),\n",
    "                low_rew_p=fixed(.1),\n",
    "                p_stay=fixed(.95),\n",
    "                T=fixed(200),\n",
    "                threshold=(.0, 1., .001))\n",
    "\n",
    "widget.children[-2].description='Run Simulation'\n",
    "widget.children[-2].style.button_color='lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 5: Implementing Value function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.147609Z",
     "iopub.status.busy": "2021-05-25T01:18:17.147069Z",
     "iopub.status.idle": "2021-05-25T01:18:17.191327Z",
     "shell.execute_reply": "2021-05-25T01:18:17.191823Z"
    },
    "outputId": "4f8ce4d5-6e01-4306-a951-6c2f8e8270d5"
   },
   "outputs": [],
   "source": [
    "#@title Video 5: Evaluate policy\n",
    "video = YouTubeVideo(id='aJhffROC74w', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 5\n",
    "Now we have generated behavior for a policy parameterized by a threshold. While it seems clear that this is at least better than being lazy, we want to know how good it is. For that, we will calculate a _value function_. We will use this value to compare different policies, and maximize the amount of fish we catch while minimizing our effort. \n",
    "\n",
    "Specifically, here the value is total expected utility per unit time.\n",
    "\n",
    "$$V(\\theta) = \\frac{1}{T}\\left(\\sum_t U_s(s_t) + U_a(a_t)\\right)$$ \n",
    "\n",
    "where $U_s(s_t)$ is the instantaneous utility (reward) from the site, and $U_a(a_t)$ is the utility (negative cost) for the chosen action. Here, the action cost is 0 if you stay, and `cost_sw` if you switch. \n",
    "\n",
    "We could take this average mathematically over the probabilities of rewards and actions. More simply, we get the same answer by simply averaging the _actual_ rewards and costs over a long time, so that's what you should do.\n",
    "\n",
    "\n",
    "**Instructions**\n",
    "* Fill in the function `value_function(measurement, act, cost_sw)` given a sequence of measurements, actions, and the cost of switching.\n",
    "* Visually find the threshold that yields the highest total value. We have provided code for plotting value versus threshold. The threshold $\\theta^*$ with the highest value gives the optimal policy for controlling where you should fish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.197447Z",
     "iopub.status.busy": "2021-05-25T01:18:17.196215Z",
     "iopub.status.idle": "2021-05-25T01:18:17.198049Z",
     "shell.execute_reply": "2021-05-25T01:18:17.198495Z"
    }
   },
   "outputs": [],
   "source": [
    "def value_function(measurement, act, cost_sw):\n",
    "  \"\"\"\n",
    "  value function\n",
    "\n",
    "  Args:\n",
    "    act (numpy array of string): length T with each element\n",
    "                                  taking value \"stay\" or \"switch\"\n",
    "    cost_sw (float): the cost of switching side\n",
    "    measurement (numpy array of binary): whether a reward is obtained\n",
    "\n",
    "  Returns:\n",
    "    value (float): expected utility per unit time\n",
    "  \"\"\"\n",
    "  act_int = (act == \"switch\").astype(int)\n",
    "  T = len(measurement)\n",
    "  ############################################################################\n",
    "  ## Insert your code here to:\n",
    "  ##        compute the value function = rate of catching fish - costs\n",
    "  ##\n",
    "  ## complete the function and remove\n",
    "  raise NotImplementedError(\"Student exercise: Please complete <value>\")\n",
    "  ############################################################################\n",
    "  # Calculate the value function\n",
    "  value = ...\n",
    "\n",
    "  return value\n",
    "\n",
    "# uncomment the line below to test your function\n",
    "# test_value_function()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.204296Z",
     "iopub.status.busy": "2021-05-25T01:18:17.203726Z",
     "iopub.status.idle": "2021-05-25T01:18:17.206324Z",
     "shell.execute_reply": "2021-05-25T01:18:17.206779Z"
    },
    "outputId": "94f5c068-5642-4122-d138-2d220ab871a4"
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "def value_function(measurement, act, cost_sw):\n",
    "  \"\"\"\n",
    "  value function\n",
    "\n",
    "  Args:\n",
    "    act (numpy array of string): length T with each element\n",
    "                                  taking value \"stay\" or \"switch\"\n",
    "    cost_sw (float): the cost of switching side\n",
    "    measurement (numpy array of binary): whether a reward is obtained\n",
    "\n",
    "  Returns:\n",
    "    value (float): expected utility per unit time\n",
    "  \"\"\"\n",
    "  act_int = (act == \"switch\").astype(int)\n",
    "  T = len(measurement)\n",
    "  # Calculate the value function\n",
    "  value = np.sum(measurement - act_int * cost_sw) / T\n",
    "\n",
    "  return value\n",
    "\n",
    "test_value_function()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 429
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:17.254102Z",
     "iopub.status.busy": "2021-05-25T01:18:17.227604Z",
     "iopub.status.idle": "2021-05-25T01:18:19.548228Z",
     "shell.execute_reply": "2021-05-25T01:18:19.547729Z"
    },
    "outputId": "2b9a2d36-9c28-4241-ee9f-ad7932f51a3a"
   },
   "outputs": [],
   "source": [
    "# Brute force search for optimal policy: loop over thresholds and compute value for each.\n",
    "# This function is needed for the second exercise.\n",
    "\n",
    "\n",
    "def value_threshold(params, cost_sw, step):\n",
    "  threshold_array = np.arange(0, .5 + step, step)\n",
    "  value_array = np.zeros(threshold_array.shape)\n",
    "\n",
    "  T, p_stay, high_rew_p, low_rew_p, _ = params\n",
    "\n",
    "  for i in range(len(threshold_array)):\n",
    "    threshold = threshold_array[i]\n",
    "\n",
    "    params = [T, p_stay, high_rew_p, low_rew_p, threshold]\n",
    "    binaryHMM_test = binaryHMM_belief(params, choose_policy=\"threshold\")\n",
    "    belief, loc, act, measurement, fish_state = binaryHMM_test.generate_process()\n",
    "\n",
    "    value_array[i] = value_function(measurement, act, cost_sw)\n",
    "\n",
    "  return threshold_array, value_array\n",
    "\n",
    "\n",
    "plot_value_threshold(cost_sw=0.5, p_stay=0.95, high_rew_p=0.4, low_rew_p=0.1, T=10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "---\n",
    "# Section 6 (Optional): Different task, different optimal policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:19.554015Z",
     "iopub.status.busy": "2021-05-25T01:18:19.553461Z",
     "iopub.status.idle": "2021-05-25T01:18:19.597253Z",
     "shell.execute_reply": "2021-05-25T01:18:19.596410Z"
    },
    "outputId": "3c58551d-9a5f-425e-c099-f8352f5998ab"
   },
   "outputs": [],
   "source": [
    "#@title Video 6: Sensitivity of optimal policy\n",
    "video = YouTubeVideo(id='wd8IVsKoEfA', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Exercise 6\n",
    "This is a thinking exercise.\n",
    "\n",
    "**Instructions:**\n",
    "After plotting value versus threshold, adjust various task parameters using the sliders below, and observe how the optimal threshold moves with\n",
    "* switching cost (`cost_sw`)\n",
    "* fish dynamics (`p_switch`)\n",
    "* probability of catching fish on each side, `low_rew_p` and `high_rew_p`\n",
    "\n",
    "Can you explain why the optimal threshold changes with these parameters?\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "478143ca0c044f5fac45f04d31f18af9",
      "9760739839d142daba1dc58bb9e75a30",
      "59bba74d07b845a9b253ad49174c4b7d",
      "495cc73a57964be09b28bb725d836ac4",
      "61182e80d8b04aa6bbcfd330e9ac4220",
      "ffd72f74ca274d5c8b91d26967f5098d",
      "0fdac623df8c4e7294f225c9775956b1",
      "bce2037903a1439dad640644d5efcf6b",
      "7deaa68ce2984f9c9a191d294718d9e0",
      "df231b0efaf24d8d96b6f0c277878295",
      "122a69af9cbd47d7a9b320e5a7ccc866",
      "a87e14261d4449fa948238326c2dc3d6",
      "88b84d51cac24a7eb4068f75d9c8e317",
      "0ea6aeaf5ddb45f8811d81d38ff2f71f",
      "052ac25f36ce411f86b8a23077cf914c",
      "96f76bae56fa48c196cd283600ee795c",
      "20c1ab71b6924fdf914cdd2c28c99a50",
      "128d983537494f9e9da575a6c8dcb593",
      "e28a996ca47e428eab33170271d758f4",
      "480f3f09749846c091ee6bc74271a028",
      "9db1c9c7c3e747429bb4d75877c16a54"
     ]
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:19.661335Z",
     "iopub.status.busy": "2021-05-25T01:18:19.660726Z",
     "iopub.status.idle": "2021-05-25T01:18:19.684679Z",
     "shell.execute_reply": "2021-05-25T01:18:19.672657Z"
    },
    "outputId": "fd8a8ae4-005d-458e-c0cf-356f9b93a7b6"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "widget=interactive(plot_value_threshold, {'manual': True},\n",
    "                   T=fixed(10000),\n",
    "                   p_stay=(0.5, 1., 0.001),\n",
    "                   high_rew_p=(0., 1., 0.001),\n",
    "                   low_rew_p=(0., 1., 0.001),\n",
    "                   cost_sw=(0., 2., .1),\n",
    "                   step=fixed(0.1))\n",
    "\n",
    "widget.children[-2].description='Run Simulation'\n",
    "widget.children[-2].style.button_color='lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Questions:**\n",
    "\n",
    "EXPLAIN why the optimal threshold changes for:\n",
    "* lower switching cost?\n",
    "* faster fish dynamics?\n",
    "* rarer fish?\n",
    "\n",
    "Note that it may require long simulations to see subtle changes in values of different policies, so look for coarse trends first."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Answers:**\n",
    "* High switching cost means that you should be more certain that the other side is better before committing to change sides. This means that beliefs must fall below a threshold before acting. Conversely, a lower switching cost allows you more flexibility to switch at less stringent thresholds. In the limit of _zero_ switching cost, you should always switch whenever you think the other side is better, even if it's just 51%, and even if you switch every time step.\n",
    "* Faster fish dynamics (lower `p_stay`) also promotes faster switching, because you cannot plan as far into the future. In that case you must base your decisions on more immediate evidence, but since you still pay the same switching cost that cost is a higher fraction of your predictable rewards. And thus you should be more conservative, and switch only when you are more confident.\n",
    "* When `high_rew_p` and/or `low_rew_p` decreases, your predictions become less reliable, again encouraging you to require more confidence before committing to a switch.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 518
    },
    "colab_type": "code",
    "execution": {
     "iopub.execute_input": "2021-05-25T01:18:19.689690Z",
     "iopub.status.busy": "2021-05-25T01:18:19.689135Z",
     "iopub.status.idle": "2021-05-25T01:18:19.735521Z",
     "shell.execute_reply": "2021-05-25T01:18:19.734984Z"
    },
    "outputId": "c4bfa510-9328-4e6f-a159-b475201d940d"
   },
   "outputs": [],
   "source": [
    "#@title Video 7: From discrete to continuous control\n",
    "video = YouTubeVideo(id='ndCMgdjv9Gg', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W3D3_Tutorial1",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": [
     "# Tutorial 1- Optimal Control for Discrete State\n",
     "\n",
     "Please execute the cell below to initialize the notebook environment.\n",
     "\n",
     "import numpy as np                 # import numpy\n",
     "import scipy               # import scipy\n",
     "import random                      # import basic random number generator functions\n",
     "from scipy.linalg import inv\n",
     "\n",
     "import matplotlib.pyplot as plt    # import matplotlib\n",
     "\n",
     "---\n",
     "\n",
     "## Tutorial objectives\n",
     "\n",
     "In this tutorial, we will implement a binary HMM task.\n",
     "\n",
     "---\n",
     "\n",
     "## Task Description\n",
     "\n",
     "There are two boxes. The box can be in a high-rewarding state ($s=1$), which means that a reward will be delivered with high probabilty $q_{high}$; or the box can be in low-rewarding state ($s=0$), then the reward will be delivered with low probabilty $q_{low}$.\n",
     "\n",
     "The states of the two boxes are latent. At a certain time, only one of the sites can be in high-rewarding state, and the other box will be the opposite. The states of the two boxes switches with a certain probability $p_{sw}$. \n",
     "\n",
     "![alt text](switching.png \"Title\")\n",
     "\n",
     "\n",
     "The agent may stay at one site for sometime. As the agent accumulates evidence about the state of the box on that site, it may choose to stay or switch to the other side with a switching cost $c$. The agent keeps beliefs on the states of the boxes, which is the posterior probability of the state being high-rewarding given all the past observations. Consider the belief on the state of the left box, we have \n",
     "\n",
     "$$b(s_t) = p(s_t = 1 | o_{0:t}, l_{0:t}, a_{0:t-1})$$\n",
     "\n",
     "where $o$ is the observation that whether a reward is obtained, $l$ is the location of the agent, $a$ is the action of staying ($a=0$) or switching($a=1$). \n",
     "\n",
     "Since the two boxes are completely anti-correlated, i.e. only one of the boxes is high-rewarded at a certain time, the the other one is low-rewarded, the belief on the two boxes should sum up to be 1. As a result, we only need to track the belief on one of the boxes. \n",
     "\n",
     "The policy of the agent depends on a threshold on beliefs. When the belief on the box on the other side gets higher than the threshold $\\theta$, the agent will switch to the other side. In other words, the agent will choose to switch when it is confident enough that the other side is high rewarding. \n",
     "\n",
     "The value function can be defined as the reward rate during a single trial.\n",
     "\n",
     "$$v(\\theta) = \\sum_t r_t - c\\cdot 1_{a_t = 1}$$ \n",
     "\n",
     "we would like to see the relation between the threshold and the value function. \n",
     "\n",
     "### Exercise 1: Control for binary HMM\n",
     "In this excercise, we generate the dynamics for the binary HMM task as described above. \n",
     "\n",
     "# This function is the policy based on threshold\n",
     "\n",
     "def policy(threshold, bel, loc):\n",
     "    if loc == 0:\n",
     "        if bel[1]  >= threshold:\n",
     "            act = 1\n",
     "        else:\n",
     "            act = 0\n",
     "    else:  # loc = 1\n",
     "        if bel[0] >= threshold:\n",
     "            act = 1\n",
     "        else:\n",
     "            act = 0\n",
     "\n",
     "    return act\n",
     "\n",
     "# This function generates the dynamics\n",
     "\n",
     "def generateProcess(params):\n",
     "\n",
     "    T, p_sw, q_high, q_low, cost_sw, threshold = params\n",
     "    world_state = np.zeros((2, T), int)  # value :1: good box; 0: bad box\n",
     "    loc = np.zeros(T, int)  # 0: left box               1: right box\n",
     "    obs = np.zeros(T, int)  # 0: did not get food        1: get food\n",
     "    act = np.zeros(T, int)  # 0 : stay                   1: switch and get food from the other side\n",
     "    bel = np.zeros((2, T), float)  # the probability that the left box has food,\n",
     "    # then the probability that the second box has food is 1-b\n",
     "\n",
     "\n",
     "    p = np.array([1 - p_sw, p_sw])  # transition probability to good state\n",
     "    q = np.array([q_low, q_high])\n",
     "    q_mat = np.array([[1 - q_high, q_high], [1 - q_low, q_low]])\n",
     "\n",
     "    for t in range(T):\n",
     "        if t == 0:\n",
     "            world_state[0, t] = 1    # good box\n",
     "            world_state[1, t] = 1 - world_state[0, t]\n",
     "            loc[t] = 0\n",
     "            obs[t] = 0\n",
     "            bel_0 = np.random.random(1)[0]\n",
     "            bel[:, t] = np.array([bel_0, 1-bel_0])\n",
     "\n",
     "            act[t] = policy(threshold, bel[:, t], loc[t])\n",
     "\n",
     "        else:\n",
     "            world_state[0, t] = np.random.binomial(1, p[world_state[0, t - 1]])\n",
     "            world_state[1, t] = 1 - world_state[0, t]\n",
     "\n",
     "            if act[t - 1] == 0:\n",
     "                loc[t] = loc[t - 1]\n",
     "            else:  # after weitching, open the new box, deplete if any; then wait a usualy time\n",
     "                loc[t] = 1 - loc[t - 1]\n",
     "\n",
     "            # new observation\n",
     "            obs[t] = np.random.binomial(1, q[world_state[loc[t], t-1]])\n",
     "\n",
     "            # update belief posterior, p(s[t] | obs(0-t), act(0-t-1))\n",
     "            bel_0 = (bel[0, t-1] * p_sw  + bel[1, t-1] * (1 - p_sw)) * q_mat[loc[t], obs[t]]\n",
     "            bel_1 = (bel[1, t - 1] * p_sw + bel[0, t - 1] * (1 - p_sw)) * q_mat[1-loc[t], obs[t]]\n",
     "\n",
     "            bel[0, t] = bel_0 / (bel_0 + bel_1)\n",
     "            bel[1, t] = bel_1 / (bel_0 + bel_1)\n",
     "\n",
     "            act[t] = policy(threshold, bel[:, t], loc[t])\n",
     "\n",
     "    return bel, obs, act, world_state, loc\n",
     "\n",
     "# value function \n",
     "def value_function(obs, act, cost_sw, discount):\n",
     "    T = len(obs)\n",
     "    discount_time = np.array([discount ** t for t in range(T)])\n",
     "\n",
     "    #value = (np.sum(obs) - np.sum(act) * cost_sw) / T\n",
     "    value = (np.sum(np.multiply(obs, discount_time)) - np.sum(np.multiply(act, discount_time)) * cost_sw) / T\n",
     "\n",
     "    return value\n",
     "\n",
     "def switch_int(obs, act):\n",
     "    sw_t = np.where(act == 1)[0]\n",
     "    sw_int = sw_t[1:] - sw_t[:-1]\n",
     "\n",
     "    return sw_int\n",
     "\n",
     "#Plotting \n",
     "def plot_dynamics(bel, obs, act, world_state, loc):\n",
     "    T = len(obs)\n",
     "\n",
     "    showlen = min(T, 100)\n",
     "    startT = 0\n",
     "\n",
     "    endT = startT + showlen\n",
     "    showT = range(startT, endT)\n",
     "    time_range = np.linspace(0, showlen - 1)\n",
     "\n",
     "    fig_posterior, [ax0, ax1, ax_loc, ax2, ax3] = plt.subplots(5, 1, figsize=(15, 10))\n",
     "\n",
     "    ax0.plot(world_state[0, showT], color='dodgerblue', markersize=10, linewidth=3.0)\n",
     "    ax0.set_ylabel('Left box', rotation=360, fontsize=22)\n",
     "    ax0.yaxis.set_label_coords(-0.1, 0.25)\n",
     "    ax0.set_xticks(np.arange(0, showlen, 10))\n",
     "    ax0.tick_params(axis='both', which='major', labelsize=18)\n",
     "    ax0.set_xlim([0, showlen])\n",
     "\n",
     "\n",
     "    ax3.plot(world_state[1, showT], color='dodgerblue', markersize=10, linewidth=3.0)\n",
     "    ax3.set_ylabel('Right box', rotation=360, fontsize=22)\n",
     "    ax3.yaxis.set_label_coords(-0.1, 0.25)\n",
     "    ax3.tick_params(axis='both', which='major', labelsize=18)\n",
     "    ax3.set_xlim([0, showlen])\n",
     "    ax3.set_xticks(np.arange(0, showlen, 10))\n",
     "\n",
     "    ax1.plot(bel[0, showT], color='dodgerblue', markersize=10, linewidth=3.0)\n",
     "    ax1.plot(time_range, threshold * np.ones(time_range.shape), 'r--')\n",
     "    ax1.yaxis.set_label_coords(-0.1, 0.25)\n",
     "    ax1.set_ylabel('Belief on \\n left box', rotation=360, fontsize=22)\n",
     "    ax1.tick_params(axis='both', which='major', labelsize=18)\n",
     "    ax1.set_xlim([0, showlen])\n",
     "    ax1.set_ylim([0, 1])\n",
     "    ax1.set_xticks(np.arange(0, showlen, 10))\n",
     "\n",
     "\n",
     "    ax_loc.plot(1 - loc[showT], 'g.-', markersize=12, linewidth=5, label = 'location')\n",
     "    ax_loc.plot((act[showT] - .1) * .8, 'v', markersize=10, label = 'action')\n",
     "    ax_loc.plot(obs[showT] * .5, '*', markersize=5, label = 'reward')\n",
     "    ax_loc.legend(loc=\"upper right\")\n",
     "    ax_loc.set_xlim([0, showlen])\n",
     "    ax_loc.set_ylim([0, 1])\n",
     "    #ax_loc.set_yticks([])\n",
     "    ax_loc.set_xticks([0, showlen])\n",
     "    ax_loc.tick_params(axis='both', which='major', labelsize=18)\n",
     "    labels = [item.get_text() for item in ax_loc.get_yticklabels()]\n",
     "    labels[0] = 'Right'\n",
     "    labels[-1] = 'Left'\n",
     "    ax_loc.set_yticklabels(labels)\n",
     "\n",
     "    ax2.plot(bel[1, showT], color='dodgerblue', markersize=10, linewidth=3.0)\n",
     "    ax2.plot(time_range, threshold * np.ones(time_range.shape), 'r--')\n",
     "    ax2.set_xlabel('time', fontsize=18)\n",
     "    ax2.yaxis.set_label_coords(-0.1, 0.25)\n",
     "    ax2.set_ylabel('Belief on  \\n  right box', rotation=360, fontsize=22)\n",
     "    ax2.tick_params(axis='both', which='major', labelsize=18)\n",
     "    ax2.set_xlim([0, showlen])\n",
     "    ax2.set_ylim([0, 1])\n",
     "    ax2.set_xticks(np.arange(0, showlen, 10))\n",
     "\n",
     "    plt.show()\n",
     "\n",
     "def plot_val_thre(threshold_array, value_array):\n",
     "    fig_, ax = plt.subplots(1, 1, figsize=(10, 10))\n",
     "    ax.plot(threshold_array, value_array)\n",
     "    ax.set_ylim([np.min(value_array), np.max(value_array)])\n",
     "    ax.set_title('threshold vs value')\n",
     "    ax.set_xlabel('threshold')\n",
     "    ax.set_ylabel('value')\n",
     "    plt.show()\n",
     "\n",
     "T = 5000\n",
     "p_sw = .95          # state transiton probability\n",
     "q_high = .7\n",
     "q_low = 0 #.2\n",
     "cost_sw = 1 #int(1/(1-p_sw)) - 5\n",
     "threshold = .8    # threshold of belief for switching\n",
     "discount = 1\n",
     "\n",
     "step = 0.1\n",
     "threshold_array = np.arange(0, 1 + step, step)\n",
     "value_array = np.zeros(threshold_array.shape)\n",
     "\n",
     "for i in range(len(threshold_array)):\n",
     "    threshold = threshold_array[i]\n",
     "    params = [T, p_sw, q_high, q_low, cost_sw, threshold]\n",
     "    bel, obs, act, world_state, loc = generateProcess(params)\n",
     "    value_array[i] = value_function(obs, act, cost_sw, discount)\n",
     "    sw_int = switch_int(obs, act)\n",
     "    #print(np.mean(sw_int))\n",
     "\n",
     "    if threshold == 0.8:\n",
     "        plot_dynamics(bel, obs, act, world_state, loc)\n",
     "\n",
     "plot_val_thre(threshold_array, value_array)\n",
     "\n",
     "\n",
     "\n",
     "\n",
     "\n",
     "\n",
     "\n",
     "\n"
    ]
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
